# AutoAugment, RandAugment, and AugMix for MindSpore
#
# This code implements the searched ImageNet policies with various tweaks and improvements and
# does not include any of the search code.
#
# AA and RA Implementation adapted from:
#     https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py, Apache 2.0
#
# AugMix adapted from:
#     https://github.com/google-research/augmix, Apache 2.0
#
# Hacked together by / Copyright 2020 Ross Wightman
#
# has been modified to adapt mindspore and mindformers
# ============================================================================
"""AutoAugment, RandAugment, and AugMix for MindSpore"""
import math
import random
import re

import numpy as np
import PIL
from PIL import Image, ImageOps, ImageEnhance
from mindspore.dataset.vision.transforms import PyTensorOperation

from mindformers.tools.register import MindFormerRegister, MindFormerModuleType

_PIL_VER = tuple([int(x) for x in PIL.__version__.split('.')[:2]])

_FILL = (128, 128, 128)

# This signifies the max integer that the controller RNN could predict for the
# augmentation scheme.
_MAX_LEVEL = 10.

_HPARAMS_DEFAULT = dict(
    translate_const=250,
    img_mean=_FILL,
)

_RANDOM_INTERPOLATION = (Image.BILINEAR, Image.BICUBIC)

INTERPOLATION = {'nearest': 0,
                 'linear': 2,
                 'cubic': 3}


def pil_interp(method):
    """Interpolation method selection"""
    if method.lower() == 'bicubic':
        func = Image.BICUBIC
    elif method.lower() == 'hamming':
        func = Image.HAMMING
    elif method.lower() == 'lanczos':
        func = Image.LANCZOS
    else:
        func = Image.BILINEAR
    return func


def _interpolation(kwargs):
    """_interpolation of auto_augment"""
    interpolation = kwargs.pop('resample', 'linear')
    interpolation = random.choice(interpolation) \
        if isinstance(interpolation, (list, tuple)) else interpolation
    return INTERPOLATION.get(interpolation)


def _check_args_tf(kwargs):
    """_check_args_tf of auto_augment"""
    if 'fillcolor' in kwargs and _PIL_VER < (5, 0):
        kwargs.pop('fillcolor')
    kwargs['resample'] = _interpolation(kwargs)


class PolicyHelper:
    """ auto_augment Policy helper """

    def shear_x(self, img, factor, **kwargs):
        """shear_x of auto_augment policy"""
        _check_args_tf(kwargs)
        return img.transform(img.size, Image.AFFINE, (1, factor, 0, 0, 1, 0), **kwargs)

    def shear_y(self, img, factor, **kwargs):
        """shear_y of auto_augment policy"""
        _check_args_tf(kwargs)
        return img.transform(img.size, Image.AFFINE, (1, 0, 0, factor, 1, 0), **kwargs)

    def translate_x_rel(self, img, pct, **kwargs):
        """translate_x_rel of auto_augment policy"""
        _check_args_tf(kwargs)
        return img.transform(img.size, Image.AFFINE, (1, 0, pct * img.size[0], 0, 1, 0), **kwargs)

    def translate_y_rel(self, img, pct, **kwargs):
        """translate_y_rel of auto_augment policy"""
        _check_args_tf(kwargs)
        return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pct * img.size[1]), **kwargs)

    def translate_x_abs(self, img, pixels, **kwargs):
        """translate_x_abs of auto_augment policy"""
        _check_args_tf(kwargs)
        return img.transform(img.size, Image.AFFINE, (1, 0, pixels, 0, 1, 0), **kwargs)

    def translate_y_abs(self, img, pixels, **kwargs):
        """translate_y_abs of auto_augment policy"""
        _check_args_tf(kwargs)
        return img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, pixels), **kwargs)

    def rotate(self, img, degrees, **kwargs):
        """rotate of auto_augment policy"""
        _check_args_tf(kwargs)
        if _PIL_VER >= (5, 2):
            func = img.rotate(degrees, **kwargs)
        elif _PIL_VER >= (5, 0):
            post_trans = (0, 0)
            rot_center = (img.size[0] / 2.0, img.size[1] / 2.0)
            angle = -math.radians(degrees)
            cos_val = round(math.cos(angle), 15)
            sin_val = round(math.sin(angle), 15)
            sin_val_neg = round(-math.sin(angle), 15)
            matrix = [cos_val, sin_val, 0.0, sin_val_neg, cos_val, 0.0]
            x = -rot_center[0] - post_trans[0]
            y = -rot_center[1] - post_trans[1]
            (a, b, c, d, e, f) = matrix
            matrix[2] = a * x + b * y + c
            matrix[5] = d * x + e * y + f
            matrix[2] += rot_center[0]
            matrix[5] += rot_center[1]
            func = img.transform(img.size, Image.AFFINE, matrix, **kwargs)
        else:
            func = img.rotate(degrees, resample=kwargs['resample'])
        return func

    def auto_contrast(self, img, **__):
        """auto_contrast of auto_augment policy"""
        return ImageOps.autocontrast(img)

    def invert(self, img, **__):
        """invert of auto_augment policy"""
        return ImageOps.invert(img)

    def equalize(self, img, **__):
        """equalize of auto_augment policy"""
        return ImageOps.equalize(img)

    def solarize(self, img, thresh, **__):
        """solarize of auto_augment policy"""
        return ImageOps.solarize(img, thresh)

    def solarize_add(self, img, add, thresh=128, **__):
        """solarize_add of auto_augment policy"""
        lut = [(min(255, i + add) if i < thresh else i) for i in range(256)]
        if img.mode == "RGB" and len(lut) == 256:
            lut = lut + lut + lut
        func = img.point(lut) if img.mode in ("L", "RGB") else img
        return func

    def posterize(self, img, bits_to_keep, **__):
        """posterize of auto_augment policy"""
        func = img if bits_to_keep >= 8 else ImageOps.posterize(img, bits_to_keep)
        return func

    def contrast(self, img, factor, **__):
        """contrast of auto_augment policy"""
        return ImageEnhance.Contrast(img).enhance(factor)

    def color(self, img, factor, **__):
        """color of auto_augment policy"""
        return ImageEnhance.Color(img).enhance(factor)

    def brightness(self, img, factor, **__):
        """brightness of auto_augment policy"""
        return ImageEnhance.Brightness(img).enhance(factor)

    def sharpness(self, img, factor, **__):
        """sharpness of auto_augment policy"""
        return ImageEnhance.Sharpness(img).enhance(factor)


class LevelHelper:
    """Level helper """

    def _randomly_negate(self, v):
        """With 50% prob, negate the value"""
        return -v if random.random() > 0.5 else v

    def rotate_level_to_arg(self, level, _):
        """_randomly_negate of auto_augment level"""
        # range [-30, 30]
        level = self._randomly_negate((level / _MAX_LEVEL) * 30.)
        return (level,)

    def enhance_level_to_arg(self, level, _):
        """enhance_level_to_arg of auto_augment level"""
        # range [0.1, 1.9]
        return ((level / _MAX_LEVEL) * 1.8 + 0.1,)

    def enhance_increasing_level_to_arg(self, level, _):
        """enhance_increasing_level_to_arg of auto_augment level"""
        # the 'no change' level is 1.0, moving away from that towards 0. or 2.0 increases the enhancement blend
        # range [0.1, 1.9]
        level = 1.0 + self._randomly_negate((level / _MAX_LEVEL) * .9)
        return (level,)

    def shear_level_to_arg(self, level, _):
        """shear_level_to_arg of auto_augment level"""
        # range [-0.3, 0.3]
        level = self._randomly_negate((level / _MAX_LEVEL) * 0.3)
        return (level,)

    def translate_abs_level_to_arg(self, level, hparams):
        """translate_abs_level_to_arg of auto_augment level"""
        translate_const = hparams['translate_const']
        level = self._randomly_negate((level / _MAX_LEVEL) * float(translate_const))
        return (level,)

    def translate_rel_level_to_arg(self, level, hparams):
        """translate_rel_level_to_arg of auto_augment level"""
        # default range [-0.45, 0.45]
        translate_pct = hparams.get('translate_pct', 0.45)
        level = self._randomly_negate((level / _MAX_LEVEL) * translate_pct)
        return (level,)

    def posterize_level_to_arg(self, level, _):
        """posterize_level_to_arg of auto_augment level"""
        # As per Tensorflow TPU EfficientNet impl
        # range [0, 4], 'keep 0 up to 4 MSB of original image'
        # intensity/severity of augmentation decreases with level
        return (int((level / _MAX_LEVEL) * 4),)

    def posterize_increasing_level_to_arg(self, level, hparams):
        """posterize_increasing_level_to_arg of auto_augment level"""
        # As per Tensorflow models research and UDA impl
        # range [4, 0], 'keep 4 down to 0 MSB of original image',
        # intensity/severity of augmentation increases with level
        return (4 - self.posterize_level_to_arg(level, hparams)[0],)

    def posterize_original_level_to_arg(self, level, _):
        """posterize_original_level_to_arg of auto_augment level"""
        # As per original AutoAugment paper description
        # range [4, 8], 'keep 4 up to 8 MSB of image'
        # intensity/severity of augmentation decreases with level
        return (int((level / _MAX_LEVEL) * 4) + 4,)

    def solarize_level_to_arg(self, level, _):
        """solarize_level_to_arg of auto_augment level"""
        # range [0, 256]
        # intensity/severity of augmentation decreases with
        return (int((level / _MAX_LEVEL) * 256),)

    def solarize_increasing_level_to_arg(self, level, hparams):
        """solarize_increasing_level_to_arg of auto_augment level"""
        # range [0, 256]
        # intensity/severity of augmentation increases with level
        return (256 - self.solarize_level_to_arg(level, hparams)[0],)

    def solarize_add_level_to_arg(self, level, _):
        """solarize_add_level_to_arg of auto_augment level"""
        # range [0, 110]
        return (int((level / _MAX_LEVEL) * 110),)


lh = LevelHelper()

LEVEL_TO_ARG = {
    'AutoContrast': None,
    'Equalize': None,
    'Invert': None,
    'Rotate': lh.rotate_level_to_arg,
    # There are several variations of the posterize level scaling in various Tensorflow/Google repositories/papers
    'Posterize': lh.posterize_level_to_arg,
    'PosterizeIncreasing': lh.posterize_increasing_level_to_arg,
    'PosterizeOriginal': lh.posterize_original_level_to_arg,
    'Solarize': lh.solarize_level_to_arg,
    'SolarizeIncreasing': lh.solarize_increasing_level_to_arg,
    'SolarizeAdd': lh.solarize_add_level_to_arg,
    'Color': lh.enhance_level_to_arg,
    'ColorIncreasing': lh.enhance_increasing_level_to_arg,
    'Contrast': lh.enhance_level_to_arg,
    'ContrastIncreasing': lh.enhance_increasing_level_to_arg,
    'Brightness': lh.enhance_level_to_arg,
    'BrightnessIncreasing': lh.enhance_increasing_level_to_arg,
    'Sharpness': lh.enhance_level_to_arg,
    'SharpnessIncreasing': lh.enhance_increasing_level_to_arg,
    'ShearX': lh.shear_level_to_arg,
    'ShearY': lh.shear_level_to_arg,
    'TranslateX': lh.translate_abs_level_to_arg,
    'TranslateY': lh.translate_abs_level_to_arg,
    'TranslateXRel': lh.translate_rel_level_to_arg,
    'TranslateYRel': lh.translate_rel_level_to_arg,
}

ph = PolicyHelper()

NAME_TO_OP = {
    'AutoContrast': ph.auto_contrast,
    'Equalize': ph.equalize,
    'Invert': ph.invert,
    'Rotate': ph.rotate,
    'Posterize': ph.posterize,
    'PosterizeIncreasing': ph.posterize,
    'PosterizeOriginal': ph.posterize,
    'Solarize': ph.solarize,
    'SolarizeIncreasing': ph.solarize,
    'SolarizeAdd': ph.solarize_add,
    'Color': ph.color,
    'ColorIncreasing': ph.color,
    'Contrast': ph.contrast,
    'ContrastIncreasing': ph.contrast,
    'Brightness': ph.brightness,
    'BrightnessIncreasing': ph.brightness,
    'Sharpness': ph.sharpness,
    'SharpnessIncreasing': ph.sharpness,
    'ShearX': ph.shear_x,
    'ShearY': ph.shear_y,
    'TranslateX': ph.translate_x_abs,
    'TranslateY': ph.translate_y_abs,
    'TranslateXRel': ph.translate_x_rel,
    'TranslateYRel': ph.translate_y_rel,
}


class AugmentOp:
    """AugmentOp"""

    def __init__(self, name, prob=0.5, magnitude=10, hparams=None):
        hparams = hparams or _HPARAMS_DEFAULT
        self.aug_fn = NAME_TO_OP[name]
        self.level_fn = LEVEL_TO_ARG[name]
        self.prob = prob
        self.magnitude = magnitude
        self.hparams = hparams.copy()
        self.kwargs = dict(
            fillcolor=tuple(hparams['img_mean']) if 'img_mean' in hparams else _FILL,
            resample=hparams['interpolation'] if 'interpolation' in hparams else _RANDOM_INTERPOLATION,
        )

        # If magnitude_std is > 0, we introduce some randomness
        # in the usually fixed policy and sample magnitude from a normal distribution
        # with mean `magnitude` and std-dev of `magnitude_std`.
        # NOTE This is my own hack, being tested, not in papers or reference impls.
        # If magnitude_std is inf, we sample magnitude from a uniform distribution
        self.magnitude_std = self.hparams.get('magnitude_std', 0)

    def __call__(self, img):
        """apply augment"""
        if self.prob < 1.0 and random.random() > self.prob:
            return img
        magnitude = self.magnitude
        if self.magnitude_std:
            if self.magnitude_std == float('inf'):
                magnitude = random.uniform(0, magnitude)
            elif self.magnitude_std > 0:
                magnitude = random.gauss(magnitude, self.magnitude_std)
        magnitude = min(_MAX_LEVEL, max(0, magnitude))  # clip to valid range
        if self.level_fn is not None:
            level_args = self.level_fn(magnitude, self.hparams)
        else:
            level_args = tuple()
        return self.aug_fn(img, *level_args, **self.kwargs)


def auto_augment_policy_v0(hparams):
    """auto_augment_policy_v0"""
    # ImageNet v0 policy from TPU EfficientNet impl, cannot find a paper reference.
    policy = [
        [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
        [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
        [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
        [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
        [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
        [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
        [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
        [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
        [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
        [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
        [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
        [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
        [('Posterize', 0.4, 6), ('AutoContrast', 0.4, 7)],
        [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
        [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
        [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
        [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
        [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
        [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
        [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
        [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
        [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
        [('Posterize', 0.8, 2), ('Solarize', 0.6, 10)],  # This results in black image with Tpu posterize
        [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
        [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
    ]
    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
    return pc


def auto_augment_policy_v0r(hparams):
    """auto_augment_policy_v0r"""
    # ImageNet v0 policy from TPU EfficientNet impl, with variation of Posterize used
    # in Google research implementation (number of bits discarded increases with magnitude)
    policy = [
        [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
        [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
        [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
        [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
        [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
        [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
        [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
        [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
        [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
        [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
        [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
        [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
        [('PosterizeIncreasing', 0.4, 6), ('AutoContrast', 0.4, 7)],
        [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
        [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
        [('Rotate', 1.0, 7), ('TranslateYRel', 0.8, 9)],
        [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
        [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
        [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
        [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
        [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
        [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
        [('PosterizeIncreasing', 0.8, 2), ('Solarize', 0.6, 10)],
        [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
        [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
    ]
    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
    return pc


def auto_augment_policy_original(hparams):
    """auto_augment_policy_original"""
    policy = [
        [('PosterizeOriginal', 0.4, 8), ('Rotate', 0.6, 9)],
        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
        [('PosterizeOriginal', 0.6, 7), ('PosterizeOriginal', 0.6, 6)],
        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
        [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
        [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
        [('PosterizeOriginal', 0.8, 5), ('Equalize', 1.0, 2)],
        [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
        [('Equalize', 0.6, 8), ('PosterizeOriginal', 0.4, 6)],
        [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
        [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
        [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
        [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
        [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
        [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
        [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
        [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
    ]
    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
    return pc


def auto_augment_policy_originalr(hparams):
    """auto_augment_policy_originalr"""
    # ImageNet policy from https://arxiv.org/abs/1805.09501 with research posterize variation
    policy = [
        [('PosterizeIncreasing', 0.4, 8), ('Rotate', 0.6, 9)],
        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
        [('PosterizeIncreasing', 0.6, 7), ('PosterizeIncreasing', 0.6, 6)],
        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
        [('Equalize', 0.4, 4), ('Rotate', 0.8, 8)],
        [('Solarize', 0.6, 3), ('Equalize', 0.6, 7)],
        [('PosterizeIncreasing', 0.8, 5), ('Equalize', 1.0, 2)],
        [('Rotate', 0.2, 3), ('Solarize', 0.6, 8)],
        [('Equalize', 0.6, 8), ('PosterizeIncreasing', 0.4, 6)],
        [('Rotate', 0.8, 8), ('Color', 0.4, 0)],
        [('Rotate', 0.4, 9), ('Equalize', 0.6, 2)],
        [('Equalize', 0.0, 7), ('Equalize', 0.8, 8)],
        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
        [('Rotate', 0.8, 8), ('Color', 1.0, 2)],
        [('Color', 0.8, 8), ('Solarize', 0.8, 7)],
        [('Sharpness', 0.4, 7), ('Invert', 0.6, 8)],
        [('ShearX', 0.6, 5), ('Equalize', 1.0, 9)],
        [('Color', 0.4, 0), ('Equalize', 0.6, 3)],
        [('Equalize', 0.4, 7), ('Solarize', 0.2, 4)],
        [('Solarize', 0.6, 5), ('AutoContrast', 0.6, 5)],
        [('Invert', 0.6, 4), ('Equalize', 1.0, 8)],
        [('Color', 0.6, 4), ('Contrast', 1.0, 8)],
        [('Equalize', 0.8, 8), ('Equalize', 0.6, 3)],
    ]
    pc = [[AugmentOp(*a, hparams=hparams) for a in sp] for sp in policy]
    return pc


def auto_augment_policy(name='v0', hparams=None):
    """auto_augment_policy"""
    hparams = hparams or _HPARAMS_DEFAULT
    if name == 'original':
        func = auto_augment_policy_original(hparams)
    elif name == 'originalr':
        func = auto_augment_policy_originalr(hparams)
    elif name == 'v0':
        func = auto_augment_policy_v0(hparams)
    elif name == 'v0r':
        func = auto_augment_policy_v0r(hparams)
    else:
        raise ValueError(f'Unknown AA policy {name}')
    return func


class AutoAugment:
    """AutoAugment"""

    def __init__(self, policy):
        self.policy = policy

    def __call__(self, img):
        """apply autoaugment"""
        sub_policy = random.choice(self.policy)
        for op in sub_policy:
            img = op(img)
        return img


@MindFormerRegister.register(MindFormerModuleType.TRANSFORMS)
def auto_augment_transform(config_str, hparams):
    """
    Create a AutoAugment transform

    :param config_str: String defining configuration of auto augmentation. Consists of multiple sections separated by
    dashes ('-'). The first section defines the AutoAugment policy (one of 'v0', 'v0r', 'original', 'originalr').
    The remaining sections, not order specific determine
        'mstd' -  float std deviation of magnitude noise applied
    Ex 'original-mstd0.5' results in AutoAugment with original policy, magnitude_std 0.5

    :param hparams: Other hparams (kwargs) for the AutoAugmentation scheme

    :return: A MindSpore compatible Transform
    """
    config = config_str.split('-')
    policy_name = config[0]
    config = config[1:]
    for c in config:
        cs = re.split(r'(\d.*)', c)
        if len(cs) >= 2:
            key, val = cs[:2]
            if key == 'mstd':
                # noise param injected via hparams for now
                hparams.setdefault('magnitude_std', float(val))
            else:
                raise ValueError('Unknown AutoAugment config section')
    return AutoAugment(auto_augment_policy(policy_name, hparams=hparams))


_RAND_TRANSFORMS = [
    'AutoContrast',
    'Equalize',
    'Invert',
    'Rotate',
    'Posterize',
    'Solarize',
    'SolarizeAdd',
    'Color',
    'Contrast',
    'Brightness',
    'Sharpness',
    'ShearX',
    'ShearY',
    'TranslateXRel',
    'TranslateYRel',
    # 'Cutout'  # NOTE I've implement this as random erasing separately
]

_RAND_INCREASING_TRANSFORMS = [
    'AutoContrast',
    'Equalize',
    'Invert',
    'Rotate',
    'PosterizeIncreasing',
    'SolarizeIncreasing',
    'SolarizeAdd',
    'ColorIncreasing',
    'ContrastIncreasing',
    'BrightnessIncreasing',
    'SharpnessIncreasing',
    'ShearX',
    'ShearY',
    'TranslateXRel',
    'TranslateYRel',
    # 'Cutout'  # NOTE I've implement this as random erasing separately
]

# These experimental weights are based loosely on the relative improvements mentioned in paper.
# They may not result in increased performance, but could likely be tuned to so.
_RAND_CHOICE_WEIGHTS_0 = {
    'Rotate': 0.3,
    'ShearX': 0.2,
    'ShearY': 0.2,
    'TranslateXRel': 0.1,
    'TranslateYRel': 0.1,
    'Color': .025,
    'Sharpness': 0.025,
    'AutoContrast': 0.025,
    'Solarize': .005,
    'SolarizeAdd': .005,
    'Contrast': .005,
    'Brightness': .005,
    'Equalize': .005,
    'Posterize': 0,
    'Invert': 0,
}


def _select_rand_weights(weight_idx=0, transforms=None):
    """_select_rand_weights"""
    transforms = transforms or _RAND_TRANSFORMS
    if weight_idx != 0:
        raise ValueError(f"only one set of weights currently, get weight_idx {weight_idx}")
    rand_weights = _RAND_CHOICE_WEIGHTS_0
    probs = [rand_weights[k] for k in transforms]
    probs /= np.sum(probs)
    return probs


def rand_augment_ops(magnitude=10, hparams=None, transforms=None):
    """rand_augment_ops"""
    hparams = hparams or _HPARAMS_DEFAULT
    transforms = transforms or _RAND_TRANSFORMS
    return [AugmentOp(
        name, prob=0.5, magnitude=magnitude, hparams=hparams) for name in transforms]


class RandAugment(PyTensorOperation):
    """RandAugment"""

    def __init__(self, ops, num_layers=2, choice_weights=None):
        super(RandAugment, self).__init__()
        self.ops = ops
        self.num_layers = num_layers
        self.choice_weights = choice_weights

    def __call__(self, img):
        # no replacement when using weighted choice
        ops = np.random.choice(
            self.ops, self.num_layers, replace=self.choice_weights is None, p=self.choice_weights)
        for op in ops:
            img = op(img)
        return img


@MindFormerRegister.register(MindFormerModuleType.TRANSFORMS)
def rand_augment_transform(config_str, hparams):
    """
    Create a RandAugment transform

    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
    dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
    sections, not order specific determine
        'm' - integer magnitude of rand augment
        'n' - integer num layers (number of transform ops selected per image)
        'w' - integer probabiliy weight index (index of a set of weights to influence choice of op)
        'mstd' -  float std deviation of magnitude noise applied
        'inc' - integer (bool), use augmentations that increase in severity with magnitude (default: 0)
    Ex 'rand-m9-n3-mstd0.5' results in RandAugment with magnitude 9, num_layers 3, magnitude_std 0.5
    'rand-mstd1-w0' results in magnitude_std 1.0, weights 0, default magnitude of 10 and num_layers 2

    :param hparams: Other hparams (kwargs) for the RandAugmentation scheme

    :return: A MindSpore compatible Transform
    """
    magnitude = _MAX_LEVEL  # default to _MAX_LEVEL for magnitude (currently 10)
    num_layers = 2  # default to 2 ops per image
    weight_idx = None  # default to no probability weights for op choice
    transforms = _RAND_TRANSFORMS
    config = config_str.split('-')
    if config[0] != 'rand':
        raise ValueError(f"config[0] is not rand, get {config[0]}")
    # [rand, m9, mstd0.5, inc1]
    config = config[1:]
    for c in config:
        cs = re.split(r'(\d.*)', c)
        if len(cs) >= 2:
            key, val = cs[:2]
            if key == 'mstd':
                # noise param injected via hparams for now
                hparams.setdefault('magnitude_std', float(val))
            elif key == 'inc':
                if bool(val):
                    transforms = _RAND_INCREASING_TRANSFORMS
            elif key == 'm':
                magnitude = int(val)
            elif key == 'n':
                num_layers = int(val)
            elif key == 'w':
                weight_idx = int(val)
            else:
                raise ValueError('Unknown RandAugment config section')
    ra_ops = rand_augment_ops(magnitude=magnitude, hparams=hparams, transforms=transforms)
    choice_weights = None if weight_idx is None else _select_rand_weights(weight_idx)
    return RandAugment(ra_ops, num_layers, choice_weights=choice_weights)


_AUGMIX_TRANSFORMS = [
    'AutoContrast',
    'ColorIncreasing',  # not in paper
    'ContrastIncreasing',  # not in paper
    'BrightnessIncreasing',  # not in paper
    'SharpnessIncreasing',  # not in paper
    'Equalize',
    'Rotate',
    'PosterizeIncreasing',
    'SolarizeIncreasing',
    'ShearX',
    'ShearY',
    'TranslateXRel',
    'TranslateYRel',
]


def augmix_ops(magnitude=10, hparams=None, transforms=None):
    """augmix_ops"""
    hparams = hparams or _HPARAMS_DEFAULT
    transforms = transforms or _AUGMIX_TRANSFORMS
    return [AugmentOp(
        name, prob=1.0, magnitude=magnitude, hparams=hparams) for name in transforms]


class AugMixAugment:
    """ AugMix Transform
    Adapted and improved from impl here
    From paper: 'AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty -
    https://arxiv.org/abs/1912.02781
    """

    def __init__(self, ops, alpha=1., width=3, depth=-1, blended=False):
        self.ops = ops
        self.alpha = alpha
        self.width = width
        self.depth = depth
        self.blended = blended  # blended mode is faster but not well tested

    def _calc_blended_weights(self, ws, m):
        """_calc_blended_weights of AugMixAugment"""
        ws = ws * m
        cump = 1.
        rws = []
        for w in ws[::-1]:
            alpha = w / cump
            cump *= (1 - alpha)
            rws.append(alpha)
        return np.array(rws[::-1], dtype=np.float32)

    def _apply_blended(self, img, mixing_weights, m):
        """_apply_blended of AugMixAugment"""
        # This is my first crack and implementing a slightly faster mixed augmentation. Instead
        # of accumulating the mix for each chain in a Numpy array and then blending with original,
        # it recomputes the blending coefficients and applies one PIL image blend per chain.
        # TODO the results appear in the right ballpark but they differ by more than rounding.
        img_orig = img.copy()
        ws = self._calc_blended_weights(mixing_weights, m)
        for w in ws:
            depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
            ops = np.random.choice(self.ops, depth, replace=True)
            img_aug = img_orig  # no ops are in-place, deep copy not necessary
            for op in ops:
                img_aug = op(img_aug)
            img = Image.blend(img, img_aug, w)
        return img

    def _apply_basic(self, img, mixing_weights, m):
        """_apply_basic of AugMixAugment"""
        # This is a literal adaptation of the paper/official implementation without normalizations and
        # PIL <-> Numpy conversions between every op. It is still quite CPU compute heavy compared to the
        # typical augmentation transforms, could use a GPU / Kornia implementation.
        img_shape = img.size[0], img.size[1], len(img.getbands())
        mixed = np.zeros(img_shape, dtype=np.float32)
        for mw in mixing_weights:
            depth = self.depth if self.depth > 0 else np.random.randint(1, 4)
            ops = np.random.choice(self.ops, depth, replace=True)
            img_aug = img  # no ops are in-place, deep copy not necessary
            for op in ops:
                img_aug = op(img_aug)
            mixed += mw * np.asarray(img_aug, dtype=np.float32)
        np.clip(mixed, 0, 255., out=mixed)
        mixed = Image.fromarray(mixed.astype(np.uint8))
        return Image.blend(img, mixed, m)

    def __call__(self, img):
        """AugMixAugment apply"""
        mixing_weights = np.float32(np.random.dirichlet([self.alpha] * self.width))
        m = np.float32(np.random.beta(self.alpha, self.alpha))
        if self.blended:
            mixed = self._apply_blended(img, mixing_weights, m)
        else:
            mixed = self._apply_basic(img, mixing_weights, m)
        return mixed


@MindFormerRegister.register(MindFormerModuleType.TRANSFORMS)
def augment_and_mix_transform(config_str, hparams):
    """ Create AugMix MindSpore transform

    :param config_str: String defining configuration of random augmentation. Consists of multiple sections separated by
    dashes ('-'). The first section defines the specific variant of rand augment (currently only 'rand'). The remaining
    sections, not order specific determine
        'm' - integer magnitude (severity) of augmentation mix (default: 3)
        'w' - integer width of augmentation chain (default: 3)
        'd' - integer depth of augmentation chain (-1 is random [1, 3], default: -1)
        'b' - integer (bool), blend each branch of chain into end result without a final blend, less CPU (default: 0)
        'mstd' -  float std deviation of magnitude noise applied (default: 0)
    Ex 'augmix-m5-w4-d2' results in AugMix with severity 5, chain width 4, chain depth 2

    :param hparams: Other hparams (kwargs) for the Augmentation transforms

    :return: A MindSpore compatible Transform
    """
    magnitude = 3
    width = 3
    depth = -1
    alpha = 1.
    blended = False
    hparams['magnitude_std'] = float('inf')
    config = config_str.split('-')
    if config[0] != 'augmix':
        raise ValueError(f"config[0] is not augmix, get {config[0]}")
    config = config[1:]
    for c in config:
        cs = re.split(r'(\d.*)', c)
        if len(cs) >= 2:
            key, val = cs[:2]
            if key == 'mstd':
                # noise param injected via hparams for now
                hparams.setdefault('magnitude_std', float(val))
            elif key == 'm':
                magnitude = int(val)
            elif key == 'w':
                width = int(val)
            elif key == 'd':
                depth = int(val)
            elif key == 'a':
                alpha = float(val)
            elif key == 'b':
                blended = bool(val)
            else:
                raise ValueError('Unknown AugMix config section')
    ops = augmix_ops(magnitude=magnitude, hparams=hparams)
    return AugMixAugment(ops, alpha=alpha, width=width, depth=depth, blended=blended)
